{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f833faf5-ddab-4355-aa7f-010d8fc3fe84",
   "metadata": {},
   "source": [
    "Lost in the middle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5e6bca1-fe13-4aa5-b224-1d74feaf8482",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install sentence-transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b4e1f7f2-865d-4188-9a3c-e54cb40848d3",
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'HuggingFaceBgeEmbeddings' from 'langchain_huggingface' (D:\\py\\ai\\langchain-env\\Lib\\site-packages\\langchain_huggingface\\__init__.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[1], line 5\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mchains\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m LLMChain,StuffDocumentsChain\n\u001b[0;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdocument_transformers\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[0;32m      3\u001b[0m     LongContextReorder\n\u001b[0;32m      4\u001b[0m )\n\u001b[1;32m----> 5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mlangchain_huggingface\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m HuggingFaceBgeEmbeddings\n\u001b[0;32m      6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mlangchain\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mvectorstores\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m Chroma\n\u001b[0;32m      8\u001b[0m \u001b[38;5;66;03m#使用huggingface托管的开源LLM来做嵌入，MiniLM-L6-v2是一个较小的LLM \u001b[39;00m\n",
      "\u001b[1;31mImportError\u001b[0m: cannot import name 'HuggingFaceBgeEmbeddings' from 'langchain_huggingface' (D:\\py\\ai\\langchain-env\\Lib\\site-packages\\langchain_huggingface\\__init__.py)"
     ]
    }
   ],
   "source": [
    "from langchain.chains import LLMChain,StuffDocumentsChain\n",
    "from langchain.document_transformers import (\n",
    "    LongContextReorder\n",
    ")\n",
    "from langchain.embeddings import HuggingFaceBgeEmbeddings\n",
    "from langchain.vectorstores import Chroma\n",
    "\n",
    "#使用huggingface托管的开源LLM来做嵌入，MiniLM-L6-v2是一个较小的LLM \n",
    "embedings = HuggingFaceBgeEmbeddings(model_name=\"all-MiniLM-L6-v2\")\n",
    "\n",
    "text = [\n",
    "    \"篮球是一项伟大的运动。\",\n",
    "    \"带我飞往月球是我最喜欢的歌曲之一。\",\n",
    "    \"凯尔特人队是我最喜欢的球队。\",\n",
    "    \"这是一篇关于波士顿凯尔特人的文件。\",\n",
    "    \"我非常喜欢去看电影。\",\n",
    "    \"波士顿凯尔特人队以20分的优势赢得了比赛。\",\n",
    "    \"这只是一段随机的文字。\",\n",
    "    \"《艾尔登之环》是过去15年最好的游戏之一。\",\n",
    "    \"L.科内特是凯尔特人队最好的球员之一。\",\n",
    "    \"拉里.伯德是一位标志性的NBA球员\"\n",
    "]\n",
    "\n",
    "retrieval = Chroma.from_texts(text,embedings).as_retriever(\n",
    "    search_kwargs={\"k\": 10}\n",
    ")\n",
    "query = \"关于凯尔特人队你都知道什么?\"\n",
    "\n",
    "#根据相关性返回文本块\n",
    "docs = retrieval.get_relevant_documents(query)\n",
    "docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0ab4ef4-3a73-44cd-b327-9ff50f5bfef2",
   "metadata": {},
   "source": [
    "## 对检索结果进行重排序\n",
    "根据论文结果，当相关信息出现在输入上下文头尾相关性最高，中间相关性最低\n",
    "- 问题相关性越低的内容块放在中间\n",
    "- 问题相关性越高的内容块放在头尾"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed71fe51-2a94-43cb-9694-b09d5b7c4b69",
   "metadata": {},
   "outputs": [],
   "source": [
    "reordering = LongContextReorder()\n",
    "reo_docs = reordering.transform_documents(docs) # 这个函数就是将相关性高的放在头尾\n",
    "reo_docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5f7c858-4c8f-41ff-89ab-301ac9704b47",
   "metadata": {},
   "outputs": [],
   "source": [
    "#检测下这种方案的精度效果\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv(\"openai.env\")\n",
    "import os\n",
    "api_key = os.environ.get(\"OPENAI_API_KEY\")\n",
    "api_base = os.environ.get(\"OPENAI_API_BASE\")\n",
    "\n",
    "\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.llms import OpenAI\n",
    "\n",
    "#设置llm\n",
    "llm = OpenAI(\n",
    "    api_key=api_key,\n",
    "    api_base=api_base,\n",
    "    model=\"gpt-3.5-turbo-instruct\",\n",
    "    temperature=0\n",
    ")\n",
    "\n",
    "document_prompt = PromptTemplate(\n",
    "    input_variables=[\"page_content\"],template=\"{page_content}\"\n",
    ")\n",
    "\n",
    "stuff_prompt_override =\"\"\"Given this text extracts:\n",
    "----------------------------------------\n",
    "{context}\n",
    "----------------------------------------\n",
    "Please answer the following questions:\n",
    "{query}\n",
    "\"\"\"\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    template=stuff_prompt_override,\n",
    "    input_variables=[\"context\",\"query\"]\n",
    ")\n",
    "\n",
    "llm_chain = LLMChain(\n",
    "    llm=llm,\n",
    "    prompt=prompt\n",
    ")\n",
    "\n",
    "WorkChain = StuffDocumentsChain(\n",
    "    llm_chain=llm_chain,\n",
    "    document_prompt=document_prompt,\n",
    "    document_variable_name=\"context\"\n",
    ")\n",
    "\n",
    "#调用\n",
    "WorkChain.run(\n",
    "    input_documents=reo_docs,\n",
    "    query=\"我最喜欢做什么事情？\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (langchain)",
   "language": "python",
   "name": "langchain-env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
